package com.wlld.yolo;

import com.wlld.MatrixTools.Matrix;
import com.wlld.config.Config;
import com.wlld.config.RZ;
import com.wlld.entity.*;
import com.wlld.function.ReLu;
import com.wlld.function.Tanh;
import com.wlld.i.OutBack;
import com.wlld.nerveCenter.NerveManager;
import com.wlld.nerveEntity.SensoryNerve;
import com.wlld.tools.Picture;
import com.wlld.tools.RgbRegression;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.geom.Rectangle2D;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.util.*;
import java.util.List;

/**
 * @param
 * @DATA
 * @Author LiDaPeng
 * @Description
 */
public class YoloByWLLD {
    private Config config;
    private NerveManager convolutionNerveManagerR;
    private NerveManager convolutionNerveManagerG;
    private NerveManager convolutionNerveManagerB;
    private NerveManager foodNerveManger;//分类拟合
    private NerveManager positionManger;//分类拟合
    private Map<Integer, YoloBody> featureMap = new HashMap<>();
    private double secondIndex = 0;
    private double firstIndex = 0;
    private boolean isEnd = false;//是否结束
    private Map<Integer, List<BorderBody>> boxBodyMap = new HashMap<>();
    private ThreeChannelMatrix backGround;//背景
    private int studyIndex = 0;
    private int minNub = 0;

    public YoloByWLLD(Config config) throws Exception {
        this.config = config;
    }

    public YoloByWLLD() throws Exception {
    }

    public Config getConfig() {
        return config;
    }

    public void setConfig(Config config) {
        this.config = config;
    }

    public boolean isEnd() {//是否结束
        return isEnd;
    }

    public void setEnd(boolean end) {
        isEnd = end;
    }

    public double getSpeed() {//获取进度
        return firstIndex * 0.5 + secondIndex * 0.5;
    }

    public void setBackGround(ThreeChannelMatrix backGround) {
        this.backGround = backGround;
    }

    public void clear() throws Exception {
        featureMap.clear();
        minNub = 0;
    }

    private void studyPosition(List<SensoryNerve> sensoryNerves, List<Double> feature, Map<Integer, Double> E, boolean isStudy, OutBack convBack
            , int[] ignores) throws Exception {
        if (sensoryNerves.size() == feature.size()) {
            for (int i = 0; i < feature.size(); i++) {
                sensoryNerves.get(i).postMessage(1, feature.get(i), isStudy, E, convBack, ignores);
            }
        } else {
            throw new Exception("size not equals,feature size:" + feature.size() + "," +
                    "sensorySize:" + sensoryNerves.size());
        }
    }

    //绘图测试
    private void draw(String path, List<Box> positions, String url2) throws Exception {
        File file = new File(path);
        FileInputStream fileInputStream = new FileInputStream(file);
        BufferedImage image2 = ImageIO.read(fileInputStream);
        int width = image2.getWidth();
        int height = image2.getHeight();
        BufferedImage bi = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        Graphics2D g2 = (Graphics2D) bi.getGraphics();
        g2.setColor(Color.RED);
        g2.drawImage(image2, 0, 0, width, height, null);
        int size = positions.size();
        System.out.println("点数size==" + size);
        for (int j = 0; j < size; j++) {//输出
            Box box = positions.get(j);
            Rectangle2D rect = new Rectangle2D.Double(box.getY(), box.getX(), box.getySize(), box.getxSize());//声明并创建矩形对象，矩形的左上角是(20，30)，宽是300，高是40
            g2.draw(rect);
        }
        String savePath = "/Users/lidapeng/Desktop/egg/" + url2 + ".jpg";
        ImageIO.write(bi, "JPEG", new FileOutputStream(savePath));
    }

    public void look(ThreeChannelMatrix threeChannelMatrix, String path, String url) throws Exception {//检测
        int boxSize = config.getBoxSize();
        ConvBack convBack = new ConvBack();
        SensoryNerve sensoryR = convolutionNerveManagerR.getSensoryNerves().get(0);
        SensoryNerve sensoryG = convolutionNerveManagerG.getSensoryNerves().get(0);
        SensoryNerve sensoryB = convolutionNerveManagerB.getSensoryNerves().get(0);
        Matrix matrixR = threeChannelMatrix.getMatrixR();
        Matrix matrixG = threeChannelMatrix.getMatrixG();
        Matrix matrixB = threeChannelMatrix.getMatrixB();
        int x = matrixB.getX();
        int y = matrixB.getY();
        FoodBack foodBack = new FoodBack();
        List<Box> boxes = new ArrayList<>();
        List<Box> boxes2 = new ArrayList<>();
        int startSize = boxSize;//2 * boxSize
        double one = config.getHeight();
        for (int i = startSize; i <= x - startSize; i += boxSize) {
            for (int j = startSize; j <= y - startSize; j += boxSize) {
                foodBack.clear();
                Matrix matrixRofSon = matrixR.getSonOfMatrix(i, j, boxSize, boxSize);
                Matrix matrixGofSon = matrixG.getSonOfMatrix(i, j, boxSize, boxSize);
                Matrix matrixBofSon = matrixB.getSonOfMatrix(i, j, boxSize, boxSize);
                List<Double> feature = new ArrayList<>();
                sensoryR.postMatrixMessage(1, matrixRofSon, false, null, convBack);
                feature.add(convBack.getMatrix().getNumber(0, 0));
                sensoryG.postMatrixMessage(1, matrixGofSon, false, null, convBack);
                feature.add(convBack.getMatrix().getNumber(0, 0));
                sensoryB.postMatrixMessage(1, matrixBofSon, false, null, convBack);
                feature.add(convBack.getMatrix().getNumber(0, 0));
                studyPosition(foodNerveManger.getSensoryNerves(), feature, null, false, foodBack, null);
                int type = foodBack.getId();//类别
                if (type == 2) {
                    studyPosition(positionManger.getSensoryNerves(), feature, null, false, convBack, null);
                    double confidence = convBack.getConfidence();//置信度
                    Box box = new Box();
                    Box box2 = new Box();
                    box.setType(type);
                    box.setConfidence(confidence);
                    box2.setX(i);
                    box2.setY(j);
                    box2.setxSize(boxSize);
                    box2.setySize(boxSize);
                    box.setX((int) (i + convBack.getXDist() * one));
                    box.setY((int) (j + convBack.getYDist() * one));
                    box.setxSize((int) (convBack.getHeight() * one));
                    box.setySize((int) (convBack.getWidth() * one));
                    boxes.add(box);
                    boxes2.add(box2);
                }
            }
        }
        System.out.println("类别数量" + boxes.size());
        NMS nms = new NMS();
        List<Box> boxList = nms.start(boxes);
        draw(path, boxList, url);
    }

    public void studyConv() throws Exception {//卷积层学习完毕
        for (int i = 0; i < minNub; i++) {
            for (Map.Entry<Integer, YoloBody> entry : featureMap.entrySet()) {
                YoloBody yoloBody = entry.getValue();
                studyPicture(yoloBody, i);
            }
            studyBackGround(false);
        }
    }

    public void studyDnn() throws Exception {//全链接学习
        ConvBack convBack = new ConvBack();
        for (int i = 0; i < minNub; i++) {
            for (Map.Entry<Integer, YoloBody> entry : featureMap.entrySet()) {
                YoloBody yoloBody = entry.getValue();
                studyType(yoloBody, convBack, i);
            }
            studyBackGround(true);
        }
    }

    private void studyType(YoloBody yoloBody, ConvBack convBack, int index) throws Exception {
        BoxBody boxBody = yoloBody.getList().get(index);//标注
        Map<Integer, Double> PE = new HashMap<>();
        Map<Integer, Double> E = new HashMap<>();
        double one = config.getHeight();
        SensoryNerve sensoryR = convolutionNerveManagerR.getSensoryNerves().get(0);
        SensoryNerve sensoryG = convolutionNerveManagerG.getSensoryNerves().get(0);
        SensoryNerve sensoryB = convolutionNerveManagerB.getSensoryNerves().get(0);
        ThreeChannelMatrix threeChannelMatrix = yoloBody.getThreeChannelMatrix();//图像
        double width = yoloBody.getPlateWidth() / one;
        double height = yoloBody.getPlateHeight() / one;
        double xDist = boxBody.getCoreDistX() / one;
        double yDist = boxBody.getCoreDistY() / one;
        double confidence = boxBody.getConfidence();//置信度需要拟合
        PE.put(1, confidence);
        PE.put(2, width);
        PE.put(3, height);
        PE.put(4, xDist);
        PE.put(5, yDist);
        Matrix matrixR = threeChannelMatrix.getMatrixR();
        Matrix matrixG = threeChannelMatrix.getMatrixG();
        Matrix matrixB = threeChannelMatrix.getMatrixB();
        int boxSize = config.getBoxSize();//格子大小
        int boxType = yoloBody.getBoxType();//类别,需要拟合one-hot
        E.put(boxType, 1D);//设置类别
        int x = boxBody.getX();//坐标
        int y = boxBody.getY();//坐标
        Matrix matrixRofSon = matrixR.getSonOfMatrix(x, y, boxSize, boxSize);
        Matrix matrixGofSon = matrixG.getSonOfMatrix(x, y, boxSize, boxSize);
        Matrix matrixBofSon = matrixB.getSonOfMatrix(x, y, boxSize, boxSize);
        List<Double> feature = new ArrayList<>();
        sensoryR.postMatrixMessage(1, matrixRofSon, false, null, convBack);
        feature.add(convBack.getMatrix().getNumber(0, 0));
        sensoryG.postMatrixMessage(1, matrixGofSon, false, null, convBack);
        feature.add(convBack.getMatrix().getNumber(0, 0));
        sensoryB.postMatrixMessage(1, matrixBofSon, false, null, convBack);
        feature.add(convBack.getMatrix().getNumber(0, 0));
        studyPosition(foodNerveManger.getSensoryNerves(), feature, E, true, null, null);
        studyPosition(positionManger.getSensoryNerves(), feature, PE, true, null, null);
    }

    private void studyBackGround(boolean isStudyType) throws Exception {//学习背景
        int boxSize = config.getBoxSize();
        Matrix matrixR = backGround.getMatrixR();
        Matrix matrixG = backGround.getMatrixG();
        Matrix matrixB = backGround.getMatrixB();
        int y = matrixR.getY() / boxSize;
        int x = matrixR.getX() / boxSize;
        int sigma = x * y;
        if (studyIndex > (sigma - 1)) {
            studyIndex = 0;
        }
        int xr = (studyIndex / y) * boxSize;
        int yr = (studyIndex % y) * boxSize;
        Matrix matrixRofSon = matrixR.getSonOfMatrix(xr, yr, boxSize, boxSize);
        Matrix matrixGofSon = matrixG.getSonOfMatrix(xr, yr, boxSize, boxSize);
        Matrix matrixBofSon = matrixB.getSonOfMatrix(xr, yr, boxSize, boxSize);
        SensoryNerve sensoryR = convolutionNerveManagerR.getSensoryNerves().get(0);
        SensoryNerve sensoryG = convolutionNerveManagerG.getSensoryNerves().get(0);
        SensoryNerve sensoryB = convolutionNerveManagerB.getSensoryNerves().get(0);
        if (!isStudyType) {
            Matrix matrixRE = new Matrix(1, 1);
            Matrix matrixGE = new Matrix(1, 1);
            Matrix matrixBE = new Matrix(1, 1);
            sensoryR.postMatrixMessage(1, matrixRofSon, true, matrixRE, null);
            sensoryG.postMatrixMessage(1, matrixGofSon, true, matrixGE, null);
            sensoryB.postMatrixMessage(1, matrixBofSon, true, matrixBE, null);
        } else {
            Map<Integer, Double> E = new HashMap<>();
            E.put(1, 1D);
            ConvBack convBack = new ConvBack();
            List<Double> feature = new ArrayList<>();
            sensoryR.postMatrixMessage(1, matrixRofSon, false, null, convBack);
            feature.add(convBack.getMatrix().getNumber(0, 0));
            sensoryG.postMatrixMessage(1, matrixGofSon, false, null, convBack);
            feature.add(convBack.getMatrix().getNumber(0, 0));
            sensoryB.postMatrixMessage(1, matrixBofSon, false, null, convBack);
            feature.add(convBack.getMatrix().getNumber(0, 0));
            System.out.println(feature);
            studyPosition(foodNerveManger.getSensoryNerves(), feature, E, true, null, null);
        }
        studyIndex++;
    }

    private void studyPicture(YoloBody yoloBody, int index) throws Exception {
        BoxBody boxBody = yoloBody.getList().get(index);//标注
        ThreeChannelMatrix threeChannelMatrix = yoloBody.getThreeChannelMatrix();//图像
        Matrix matrixR = threeChannelMatrix.getMatrixR();
        Matrix matrixG = threeChannelMatrix.getMatrixG();
        Matrix matrixB = threeChannelMatrix.getMatrixB();
        int boxSize = config.getBoxSize();//格子大小
        double one = config.getHeight();
        int classificationNub = config.getClassificationNub() - 1;
        double e = 1D / classificationNub;
        int boxType = yoloBody.getBoxType();//类别,需要拟合one-hot
        int x = boxBody.getX();//坐标
        int y = boxBody.getY();//坐标
        Matrix matrixRE = new Matrix(1, 1);
        Matrix matrixGE = new Matrix(1, 1);
        Matrix matrixBE = new Matrix(1, 1);//(boxType - 1) * e
        matrixRE.setNub(0, 0, boxType * e);
        matrixGE.setNub(0, 0, boxType - 1 * e);
        matrixBE.setNub(0, 0, boxType - 1 * e);
        Matrix matrixRofSon = matrixR.getSonOfMatrix(x, y, boxSize, boxSize);
        Matrix matrixGofSon = matrixG.getSonOfMatrix(x, y, boxSize, boxSize);
        Matrix matrixBofSon = matrixB.getSonOfMatrix(x, y, boxSize, boxSize);
        SensoryNerve sensoryR = convolutionNerveManagerR.getSensoryNerves().get(0);
        SensoryNerve sensoryG = convolutionNerveManagerG.getSensoryNerves().get(0);
        SensoryNerve sensoryB = convolutionNerveManagerB.getSensoryNerves().get(0);
        sensoryR.postMatrixMessage(1, matrixRofSon, true, matrixRE, null);
        sensoryG.postMatrixMessage(1, matrixGofSon, true, matrixGE, null);
        sensoryB.postMatrixMessage(1, matrixBofSon, true, matrixBE, null);
    }

    public void insertFeature(YoloBody yoloBody, String url) throws Exception {
        int type = yoloBody.getBoxType();
        int size = yoloBody.getList().size();
        if (minNub == 0 || size < minNub) {
            minNub = size;
        }
        Picture picture = new Picture();
        yoloBody.setThreeChannelMatrix(picture.getThreeMatrix(url));
        featureMap.put(type, yoloBody);
    }

    public void init() throws Exception {
        int x = config.getBoxSize();
        int kernLen = config.getKernLen();
        int step = config.getStep();
        int deep = 0;
        int size = 0;
        do {
            x = (x - (kernLen - step)) / step;
            if (x > 0) {
                size = x;
                deep++;
            }
        } while (x > 0);
        int classificationNub = config.getClassificationNub();
        int senNub = size * size * 3;
        convolutionNerveManagerR = initConv(deep, step, kernLen);
        convolutionNerveManagerG = initConv(deep, step, kernLen);
        convolutionNerveManagerB = initConv(deep, step, kernLen);
        foodNerveManger = new NerveManager(senNub, config.getAllHiddenNub(), classificationNub, config.getLineDeep(), new Tanh(), false
                , config.getAllLineStudyPoint(), RZ.L1, config.getlParam());
        positionManger = new NerveManager(senNub, config.getAllHiddenNub(), 5, config.getLineDeep(), new Tanh(), false
                , config.getAllLineStudyPoint(), RZ.L1, config.getlParam());
        foodNerveManger.init(true, false, true, config.isSoftMax(), 0, 0);
        positionManger.init(true, false, false, false, 0, 0);
    }

    private NerveManager initConv(int deep, int step, int kernLen) throws Exception {
        NerveManager nerveManager = new NerveManager(1, 1,
                1, deep - 1, new ReLu(),
                true, config.getConvStudyPoint(), RZ.NOT_RZ, 0);
        nerveManager.init(true, true, false, false
                , step, kernLen);
        return nerveManager;
    }

    public NerveManager getConvolutionNerveManagerR() {
        return convolutionNerveManagerR;
    }

    public NerveManager getConvolutionNerveManagerG() {
        return convolutionNerveManagerG;
    }

    public NerveManager getConvolutionNerveManagerB() {
        return convolutionNerveManagerB;
    }

    public NerveManager getFoodNerveManger() {
        return foodNerveManger;
    }
}
